package com.pig4cloud.pigx.knowledge.service.impl;

import cn.hutool.core.lang.Pair;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.pig4cloud.pigx.common.core.util.R;
import com.pig4cloud.pigx.knowledge.dto.LlmOCRFieldDTO;
import com.pig4cloud.pigx.knowledge.entity.AiOcrConfEntity;
import com.pig4cloud.pigx.knowledge.mapper.AiOcrConfMapper;
import com.pig4cloud.pigx.knowledge.service.AiAssistantService;
import com.pig4cloud.pigx.knowledge.service.AiOcrConfService;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import com.pig4cloud.pigx.knowledge.support.util.PromptBuilder;
import dev.ai4j.openai4j.chat.ResponseFormatType;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ImageContent;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ResponseFormat;
import dev.langchain4j.model.chat.request.json.JsonBooleanSchema;
import dev.langchain4j.model.chat.request.json.JsonObjectSchema;
import dev.langchain4j.model.chat.request.json.JsonSchema;
import dev.langchain4j.model.chat.request.json.JsonStringSchema;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.Response;
import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.tuple.Triple;
import org.springframework.stereotype.Service;
import org.springframework.util.MimeTypeUtils;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

import static dev.langchain4j.data.message.SystemMessage.systemMessage;
import static dev.langchain4j.data.message.UserMessage.userMessage;
import static dev.langchain4j.model.chat.request.ResponseFormatType.JSON;

/**
 * AI OCR conf Service impl
 *
 * @author lengleng
 * @date 2024/09/10
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class AiOcrConfServiceImpl extends ServiceImpl<AiOcrConfMapper, AiOcrConfEntity> implements AiOcrConfService {

	private final ModelProvider modelProvider;

	/**
	 * 解析图像
	 * @param aiOcrConf AI OCR 会议
	 * @return {@link R }
	 */
	@SneakyThrows
	@Override
	public R parseImage(AiOcrConfEntity aiOcrConf) {
		// 查询底图设置和标记字段
		AiOcrConfEntity confEntity = baseMapper.selectById(aiOcrConf.getId());
		// 需要OCR 的字段列表
		List<LlmOCRFieldDTO> fieldDTOList = JSONUtil.parseObj(confEntity.getOcrMarked())
			.getBeanList("markerList", LlmOCRFieldDTO.class);

		// 用户输入提示词增加限制
		SystemMessage systemMessage = SystemMessage.from(PromptBuilder.render("ocr-image.st"));
		String base64 = StrUtil.removePrefix(aiOcrConf.getImageResource(), "data:image/png;base64,");
		UserMessage userMessage = UserMessage.from(ImageContent.from(base64, MimeTypeUtils.IMAGE_PNG_VALUE));

		// 这里只能使用图形的模型助手
		Pair<ChatLanguageModel, AiAssistantService> aiAssistantServicePair = modelProvider.getAiVisionAssistant();
		Response<AiMessage> lllAiOcrResult = aiAssistantServicePair.getKey().generate(systemMessage, userMessage);
		log.info("OCR 结果: {}", lllAiOcrResult.content().text());
		// 调用Chat 模型总结结果
		Triple<ChatLanguageModel, AiAssistantService, String> aiJSONAssistant = modelProvider
			.getAiJSONAssistant(aiOcrConf.getChatModelName());

		ChatResponse chatResponse = aiJSONAssistant.getLeft()
			.chat(buildChatRequest(confEntity, lllAiOcrResult, fieldDTOList, aiJSONAssistant.getRight()));
		log.info("OCR JSON: {}", chatResponse.aiMessage().text());

		return R.ok(chatResponse.aiMessage().text());
	}

	/**
	 * 构建聊天请求
	 * @param confEntity conf 实体
	 * @param lllAiOcrResult lll AI OCR 结果
	 * @param fieldDTOList 字段 dtolist
	 * @param jsonModel JSON 模型
	 * @return {@link ChatRequest }
	 */
	private ChatRequest buildChatRequest(AiOcrConfEntity confEntity, Response<AiMessage> lllAiOcrResult,
			List<LlmOCRFieldDTO> fieldDTOList, String jsonModel) {
		ChatRequest.Builder builder = ChatRequest.builder();
		List<String> fieldNames = new ArrayList<>();
		fieldNames.add("isContain");
		fieldNames.addAll(fieldDTOList.stream().map(LlmOCRFieldDTO::getName).toList());
		JsonSchema jsonSchema = JsonSchema.builder()
			.name("OCResult")
			.rootElement(JsonObjectSchema.builder().properties(new LinkedHashMap<>() {
				{
					// 是否包含
					put("isContain", JsonBooleanSchema.builder().description(confEntity.getOcrPrompt()).build());
					// 目标OCR 的字段列表
					for (LlmOCRFieldDTO llmOCRFieldDTO : fieldDTOList) {
						put(llmOCRFieldDTO.getName(),
								JsonStringSchema.builder().description(llmOCRFieldDTO.getDescription()).build());
					}
				}
			}).required(fieldNames).additionalProperties(false).build())
			.build();

		// JSON Schema
		if (ResponseFormatType.JSON_SCHEMA.name().toLowerCase().equals(jsonModel)) {
			builder
				.messages(systemMessage(PromptBuilder.render("ocr-system-json.st")),
						userMessage(lllAiOcrResult.content().text()))
				.responseFormat(ResponseFormat.builder().type(JSON).jsonSchema(jsonSchema).build());
		}

		// user message append json object
		if (ResponseFormatType.JSON_OBJECT.name().toLowerCase().equals(jsonModel)) {
			builder.messages(systemMessage(PromptBuilder.render("ocr-system-json.st")), userMessage(PromptBuilder
				.render("ocr-user-json.st",
						Map.of("userInput", lllAiOcrResult.content().text(), "jsonSchema", jsonSchema.toString()))));
		}

		return builder.build();
	}

}
